import database as db
from orm import Cwe, CweClassification, Poc

conn = db.conn

def get_cwe():
    cursor = conn.cursor()
    sql = 'select * from cwe'
    cursor.execute(sql)
    cwe_list = cursor.fetchall()
    for cwe in cwe_list:
        # 看cwe是否已经出现在数据库中了
        try:
            d = Cwe.select().where(Cwe.cwe_id == cwe[1].strip()).get()
            print('passing cwe... --> {}'.format(cwe[1].strip()))
        except Exception as e:
            print('inserting cwe... --> {}'.format(cwe[2].strip()))
            entity = Cwe(cwe_id=cwe[1].strip(), cwe_name=cwe[2].strip(), description=cwe[3].strip(), extended_description=cwe[4].strip(), url=cwe[5].strip())
            entity.save()
    return cwe_list

def get_cwe_by_cve(year):
    poc_list = Poc.select().where(Poc.year == year)
    cve_list = [poc.cve_name for poc in poc_list]

    for cve_id in cve_list:
        sql = "select * from cwe_classification where cve_id='{}'".format(cve_id)
        cursor = conn.cursor()
        cursor.execute(sql)
        classification_list = cursor.fetchall()
        for classification in classification_list:
            cve_name = classification[0]
            cwe_id = classification[1]
            print('{} : {}'.format(cve_name, cwe_id))
            # 检查cwe是否在数据库中，不在就直接抛异常停止
            try:
                cwe = Cwe.select().where(Cwe.cwe_id == cwe_id.strip()).get()
                cweClassification = CweClassification(cve_name=cve_name, cwe_id=cwe_id)
                cweClassification.save()
            except Exception as e:
                print('ERROR!!!!! CWE ID not in database.')
                print(e)
                exit(0)



if __name__ == '__main__':
    # print(Cwe.select().where(Cwe.cwe_id == 'CWE-77').get())
    # get_cwe_by_cve(2023)

    pass